package com.shujia.spark.sql

import org.apache.spark.sql.expressions.UserDefinedFunction
import org.apache.spark.sql.{DataFrame, SparkSession}

object Demo13UDF {
  def main(args: Array[String]): Unit = {
    val spark: SparkSession = SparkSession
      .builder()
      .master("local")
      .appName("hive")
      .config("spark.sql.shuffle.partitions", 1)
      .getOrCreate()
    import spark.implicits._
    import org.apache.spark.sql.functions._


    val studentDF: DataFrame = spark
      .read
      .format("csv")
      .option("sep", ",")
      .schema("id STRING,name STRING, age INT,sex STRING,clazz STRING")
      .load("data/students.txt")


    /**
     * 自定义函数
     */
    val subClazz: UserDefinedFunction = udf((clazz: String) => {
      clazz.substring(0, 2)
    })

    /**
     * 1 在DSl中使用自定义函数
     */
    studentDF
      .select(subClazz($"clazz") as "flag")
    //.show()

    /**
     * 2 在sql中使用自定义函数
     */
    studentDF.createOrReplaceTempView("student")

    //注册sql使用的自定义函数
    spark.udf.register("subClazz", subClazz)

    spark.sql(
      """
        |select subClazz(clazz) from
        |student
        |""".stripMargin).show()


  }

}
